use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation, Algorithm};
use chrono::{Utc, Duration};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::{
    error::{AppError, AppResult},
    middleware::auth::Claims,
};

pub struct JwtManager {
    encoding_key: EncodingKey,
    decoding_key: DecodingKey,
    validation: Validation,
    expiration_seconds: i64,
}

impl JwtManager {
    pub fn new(secret: &str, expiration_seconds: i64) -> Self {
        let mut validation = Validation::new(Algorithm::HS256);
        validation.validate_exp = true;
        validation.validate_nbf = false;
        
        Self {
            encoding_key: EncodingKey::from_secret(secret.as_ref()),
            decoding_key: DecodingKey::from_secret(secret.as_ref()),
            validation,
            expiration_seconds,
        }
    }

    pub fn generate_token(&self, user_id: Uuid, email: &str, role: &str) -> AppResult<String> {
        let now = Utc::now();
        let exp = now + Duration::seconds(self.expiration_seconds);

        let claims = Claims {
            sub: user_id,
            email: email.to_string(),
            role: role.to_string(),
            exp: exp.timestamp() as usize,
            iat: now.timestamp() as usize,
        };

        let header = Header::new(Algorithm::HS256);
        
        encode(&header, &claims, &self.encoding_key)
            .map_err(|e| AppError::internal(format!("Token generation failed: {}", e)))
    }

    pub fn verify_token(&self, token: &str) -> AppResult<Claims> {
        decode::<Claims>(token, &self.decoding_key, &self.validation)
            .map(|data| data.claims)
            .map_err(|e| match e.kind() {
                jsonwebtoken::errors::ErrorKind::ExpiredSignature => {
                    AppError::unauthorized("Token has expired")
                }
                jsonwebtoken::errors::ErrorKind::InvalidToken => {
                    AppError::unauthorized("Invalid token")
                }
                jsonwebtoken::errors::ErrorKind::InvalidSignature => {
                    AppError::unauthorized("Invalid token signature")
                }
                _ => AppError::unauthorized("Token validation failed"),
            })
    }

    pub fn refresh_token(&self, old_token: &str) -> AppResult<String> {
        // First verify the old token (even if expired, we want to check signature)
        let mut validation = self.validation.clone();
        validation.validate_exp = false; // Allow expired tokens for refresh
        
        let claims = decode::<Claims>(old_token, &self.decoding_key, &validation)
            .map(|data| data.claims)
            .map_err(|_| AppError::unauthorized("Invalid refresh token"))?;

        // Generate new token with same user info
        self.generate_token(claims.sub, &claims.email, &claims.role)
    }

    pub fn extract_token_from_header(auth_header: &str) -> AppResult<&str> {
        if !auth_header.starts_with("Bearer ") {
            return Err(AppError::unauthorized("Invalid authorization header format"));
        }
        
        Ok(&auth_header[7..]) // Remove "Bearer " prefix
    }

    pub fn get_token_expiration(&self) -> i64 {
        self.expiration_seconds
    }

    pub fn is_token_expired(&self, token: &str) -> bool {
        match self.verify_token(token) {
            Ok(_) => false,
            Err(e) => {
                // Check if the error is specifically about expiration
                matches!(e, AppError::Authentication(ref msg) if msg.contains("expired"))
            }
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct TokenPair {
    pub access_token: String,
    pub refresh_token: String,
    pub expires_in: i64,
    pub token_type: String,
}

impl TokenPair {
    pub fn new(access_token: String, refresh_token: String, expires_in: i64) -> Self {
        Self {
            access_token,
            refresh_token,
            expires_in,
            token_type: "Bearer".to_string(),
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct RefreshTokenClaims {
    pub sub: Uuid,
    pub token_type: String, // "refresh"
    pub exp: usize,
    pub iat: usize,
}

pub struct RefreshTokenManager {
    jwt_manager: JwtManager,
    refresh_expiration_seconds: i64,
}

impl RefreshTokenManager {
    pub fn new(secret: &str, access_expiration: i64, refresh_expiration: i64) -> Self {
        Self {
            jwt_manager: JwtManager::new(secret, access_expiration),
            refresh_expiration_seconds: refresh_expiration,
        }
    }

    pub fn generate_token_pair(&self, user_id: Uuid, email: &str, role: &str) -> AppResult<TokenPair> {
        let access_token = self.jwt_manager.generate_token(user_id, email, role)?;
        let refresh_token = self.generate_refresh_token(user_id)?;
        
        Ok(TokenPair::new(
            access_token,
            refresh_token,
            self.jwt_manager.get_token_expiration(),
        ))
    }

    pub fn generate_refresh_token(&self, user_id: Uuid) -> AppResult<String> {
        let now = Utc::now();
        let exp = now + Duration::seconds(self.refresh_expiration_seconds);

        let claims = RefreshTokenClaims {
            sub: user_id,
            token_type: "refresh".to_string(),
            exp: exp.timestamp() as usize,
            iat: now.timestamp() as usize,
        };

        let header = Header::new(Algorithm::HS256);
        
        encode(&header, &claims, &self.jwt_manager.encoding_key)
            .map_err(|e| AppError::internal(format!("Refresh token generation failed: {}", e)))
    }

    pub fn verify_refresh_token(&self, token: &str) -> AppResult<RefreshTokenClaims> {
        decode::<RefreshTokenClaims>(token, &self.jwt_manager.decoding_key, &self.jwt_manager.validation)
            .map(|data| data.claims)
            .map_err(|e| match e.kind() {
                jsonwebtoken::errors::ErrorKind::ExpiredSignature => {
                    AppError::unauthorized("Refresh token has expired")
                }
                _ => AppError::unauthorized("Invalid refresh token"),
            })
    }

    pub fn refresh_access_token(&self, refresh_token: &str, email: &str, role: &str) -> AppResult<TokenPair> {
        let refresh_claims = self.verify_refresh_token(refresh_token)?;
        
        if refresh_claims.token_type != "refresh" {
            return Err(AppError::unauthorized("Invalid token type"));
        }

        self.generate_token_pair(refresh_claims.sub, email, role)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_jwt_generation_and_verification() {
        let manager = JwtManager::new("test_secret_key_that_is_long_enough", 3600);
        let user_id = Uuid::new_v4();
        let email = "test@example.com";
        let role = "user";

        let token = manager.generate_token(user_id, email, role).unwrap();
        let claims = manager.verify_token(&token).unwrap();

        assert_eq!(claims.sub, user_id);
        assert_eq!(claims.email, email);
        assert_eq!(claims.role, role);
    }

    #[test]
    fn test_invalid_token() {
        let manager = JwtManager::new("test_secret_key_that_is_long_enough", 3600);
        let result = manager.verify_token("invalid_token");
        assert!(result.is_err());
    }

    #[test]
    fn test_extract_token_from_header() {
        let header = "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9";
        let token = JwtManager::extract_token_from_header(header).unwrap();
        assert_eq!(token, "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9");

        let invalid_header = "InvalidHeader";
        let result = JwtManager::extract_token_from_header(invalid_header);
        assert!(result.is_err());
    }

    #[test]
    fn test_refresh_token_manager() {
        let manager = RefreshTokenManager::new("test_secret_key_that_is_long_enough", 3600, 86400);
        let user_id = Uuid::new_v4();
        let email = "test@example.com";
        let role = "user";

        let token_pair = manager.generate_token_pair(user_id, email, role).unwrap();
        assert!(!token_pair.access_token.is_empty());
        assert!(!token_pair.refresh_token.is_empty());
        assert_eq!(token_pair.expires_in, 3600);
        assert_eq!(token_pair.token_type, "Bearer");

        // Verify refresh token
        let refresh_claims = manager.verify_refresh_token(&token_pair.refresh_token).unwrap();
        assert_eq!(refresh_claims.sub, user_id);
        assert_eq!(refresh_claims.token_type, "refresh");
    }
}